package main

import (
	"encoding/json"
	"fmt"
	"math"
	"os"
	"sort"
	"sync"
)

// VectorItem 向量项
type VectorItem struct {
	ID         string            `json:"id"`
	Vector     []float64         `json:"vector"`
	Metadata   map[string]interface{} `json:"metadata"`
	Content    string            `json:"content"`
	FilePath   string            `json:"file_path"`
	ChunkIndex int               `json:"chunk_index"`
	LineStart  int               `json:"line_start"`
	LineEnd    int               `json:"line_end"`
}

// VectorDatabase 简单的向量数据库实现
type VectorDatabase struct {
	items    []VectorItem
	filePath string
	mutex    sync.RWMutex
}

// NewVectorDatabase 创建新的向量数据库
func NewVectorDatabase(filePath string) (*VectorDatabase, error) {
	db := &VectorDatabase{
		items:    make([]VectorItem, 0),
		filePath: filePath,
	}
	
	// 尝试加载现有数据
	if err := db.load(); err != nil {
		fmt.Printf("⚠️  无法加载现有向量数据: %v\n", err)
		fmt.Println("📝 将创建新的向量数据库")
	}
	
	return db, nil
}

// Add 添加向量项
func (db *VectorDatabase) Add(item VectorItem) error {
	db.mutex.Lock()
	defer db.mutex.Unlock()
	
	db.items = append(db.items, item)
	return db.save()
}

// Search 向量搜索
func (db *VectorDatabase) Search(queryVector []float64, topK int, threshold float64) []SearchResult {
	db.mutex.RLock()
	defer db.mutex.RUnlock()
	
	type scoredItem struct {
		item  VectorItem
		score float64
	}
	
	var results []scoredItem
	
	for _, item := range db.items {
		score := cosineSimilarity(queryVector, item.Vector)
		if score >= threshold {
			results = append(results, scoredItem{
				item:  item,
				score: score,
			})
		}
	}
	
	// 按相似度排序
	sort.Slice(results, func(i, j int) bool {
		return results[i].score > results[j].score
	})
	
	// 限制返回数量
	if len(results) > topK {
		results = results[:topK]
	}
	
	// 转换为搜索结果
	searchResults := make([]SearchResult, len(results))
	for i, result := range results {
		searchResults[i] = SearchResult{
			FilePath:   result.item.FilePath,
			Content:    result.item.Content,
			Score:      result.score,
			LineStart:  result.item.LineStart,
			LineEnd:    result.item.LineEnd,
			ChunkIndex: result.item.ChunkIndex,
		}
	}
	
	return searchResults
}

// Clear 清空数据库
func (db *VectorDatabase) Clear() error {
	db.mutex.Lock()
	defer db.mutex.Unlock()
	
	db.items = make([]VectorItem, 0)
	return db.save()
}

// Count 获取向量数量
func (db *VectorDatabase) Count() int {
	db.mutex.RLock()
	defer db.mutex.RUnlock()
	
	return len(db.items)
}

// Close 关闭数据库
func (db *VectorDatabase) Close() error {
	return db.save()
}

// save 保存数据到文件
func (db *VectorDatabase) save() error {
	data, err := json.MarshalIndent(db.items, "", "  ")
	if err != nil {
		return err
	}
	
	return os.WriteFile(db.filePath, data, 0644)
}

// load 从文件加载数据
func (db *VectorDatabase) load() error {
	data, err := os.ReadFile(db.filePath)
	if err != nil {
		return err
	}
	
	return json.Unmarshal(data, &db.items)
}

// cosineSimilarity 计算余弦相似度
func cosineSimilarity(a, b []float64) float64 {
	if len(a) != len(b) {
		return 0
	}
	
	var dotProduct, normA, normB float64
	
	for i := 0; i < len(a); i++ {
		dotProduct += a[i] * b[i]
		normA += a[i] * a[i]
		normB += b[i] * b[i]
	}
	
	if normA == 0 || normB == 0 {
		return 0
	}
	
	return dotProduct / (math.Sqrt(normA) * math.Sqrt(normB))
}